from PyQt5.QtWidgets import QWidget, QFileDialog
from .Ui_upload import Ui_Form
from PyQt5.QtCore import pyqtSlot
import cv2
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
from ultralytics import YOLO
from qfluentwidgets import InfoBar, InfoBarPosition

class UploadInterface(QWidget, Ui_Form):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.img = None
        self.model_yolo = None
        
        self.load_models()
        
    def load_models(self):
        self.model_yolo = YOLO(model='checkpoints/best.pt')
        
    def create_success_infoBar(self):
        InfoBar.success(
            title='success',
            content="检测成功！",
            orient=Qt.Horizontal,
            isClosable=True,
            position=InfoBarPosition.TOP,
            duration=2000,
            parent=self
        )
        
    @pyqtSlot()
    def btn_upload_clicked(self):
        # 打开文件选取对话框
        open_dir = 'data/yolo/imgs/cn'
        filename, _ = QFileDialog.getOpenFileName(self, '打开', open_dir)
        filename = str(filename)
        
        if filename:
            img = cv2.imread(filename)
            self.img = img
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            QImg = QImage(
                img.data, 
                img.shape[1], 
                img.shape[0], 
                img.strides[0], 
                QImage.Format_RGB888
            )
            pixmap = QPixmap.fromImage(QImg)
            label_size = self.label_exam.size()
            scaled_pixmap = pixmap.scaled(
                label_size, 
                Qt.KeepAspectRatio, 
                Qt.SmoothTransformation
            )
            self.label_exam.setPixmap(scaled_pixmap)
        
    @pyqtSlot()
    def btn_exit_clicked(self):
        exit()
        
    @pyqtSlot()
    def btn_detect_clicked(self):
        results =  self.model_yolo.predict(
            source=self.img, 
            save=True, 
            imgsz=640,
        )
        img = results[0].plot()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        QImg = QImage(
            img.data, 
            img.shape[1], 
            img.shape[0], 
            img.strides[0], 
            QImage.Format_RGB888
        )
        pixmap = QPixmap.fromImage(QImg)
        label_size = self.label_result.size()
        scaled_pixmap = pixmap.scaled(
            label_size, 
            Qt.KeepAspectRatio, 
            Qt.SmoothTransformation
        )
        self.create_success_infoBar()
        self.label_result.setPixmap(scaled_pixmap)